#!/bin/bash
# Script to perform FIO tests on NVMe devices with flexible command-line options.
# Version F.12 (Corrected logic for '-d' flag without '-s' to allow fallback to interactive menu)
#
# Usage:
#   ./fiotest [device1 [device2 ...]] [-d] [-s <spec>] [-f 0|1] [-c <cpu_config_file|0>]
#
# Pre-requisites:
#   - fio, getopt, nvme-cli, awk installed.
#   - FIO job files: SequentialPrecondition.fio, RandomPrecondition.fio, Common.fio.
#   - A 'fio_tests.conf' file using '|' as a delimiter.
# --- Color Definitions ---
_COLOR_RED="\033[31m"; _COLOR_GREEN="\033[32m"; _COLOR_YELLOW="\033[1;33m"; _COLOR_BLUE="\033[34m"; _COLOR_CYAN="\033[36m"; _COLOR_RESET="\033[0m"
export _COLOR_RED _COLOR_GREEN _COLOR_YELLOW _COLOR_BLUE _COLOR_CYAN _COLOR_RESET
# --- FIO Test Definitions (loaded from config) ---
FIO_TEST_CONFIG_FILE="fio_tests.conf"
FIO_TEST_MENU_TEXTS=()
FIO_TEST_LOG_SUFFIXES=()
FIO_TEST_BASE_JOBS=()
FIO_TEST_RW_OPTS=()
FIO_TEST_BS_OPTS=()
FIO_TEST_IODEPTH_OPTS=()
FIO_TEST_NUMJOBS_OPTS=()
# --- Global Script Variables ---
LOGDIR="./fio_result_$(date +%Y%m%d_%H%M%S)"
mkdir -p "$LOGDIR" || { echo "Error: Failed to create log directory '$LOGDIR'. Exiting." >&2; exit 1; }
LOGFILE="$LOGDIR/script_output.log"
export LOGDIR LOGFILE
# --- Initialize option variables ---
CLI_DEVICE_ARGS=()
SELECTED_FIO_STEP_INDICES=()
PERFORM_FORMATTING=""
AUTO_SELECT_ALL_NON_SYSTEM_NVME=false
RUN_FULL_FIO_VIA_S0=false
USER_SPECIFIED_CPU_CONFIG_FILE=""
DISABLE_CPU_BINDING="false"
export DISABLE_CPU_BINDING
DEVICES=()
TEST_TYPE=""
declare -A FIO_PIDS
declare -gA DEVICE_HAS_PARTITIONS
# --- Logging Function ---
log() {
    local msg="$1"
    local color="${2:-$_COLOR_RESET}"
    local ts
    ts=$(date "+%Y-%m-%d %H:%M:%S")
    echo -e "${ts} - ${color}${msg}${_COLOR_RESET}" | tee -a "$LOGFILE"
}
export -f log
# --- #################### ---
# --- HELPER FUNCTIONS     ---
# --- #################### ---
# --- Usage Display Function ---
display_usage() {
cat << EOF
${_COLOR_CYAN}fiotest - FIO Test Automation Script for NVMe Devices${_COLOR_RESET}
${_COLOR_YELLOW}Usage:${_COLOR_RESET}
  ./fiotest [device1 [device2 ...]] [OPTIONS]
${_COLOR_YELLOW}Description:${_COLOR_RESET}
  This script automates FIO performance testing on NVMe devices. It is highly
  configurable via command-line options and the 'fio_tests.conf' file.
${_COLOR_YELLOW}Options:${_COLOR_RESET}
  ${_COLOR_GREEN}-d, --all-non-system-nvme${_COLOR_RESET}
                          Automatically select all non-system NVMe devices.
  ${_COLOR_GREEN}-s, --select-fio-step <spec>${_COLOR_RESET}
                          Select which FIO test(s) to run from 'fio_tests.conf'.
                          ${_COLOR_CYAN}<spec> = 0${_COLOR_RESET}: Run all defined tests.
                          ${_COLOR_CYAN}<spec> = <num>${_COLOR_RESET}: Run only the specified test number.
                          ${_COLOR_CYAN}<spec> = <num1,num2,...>${_COLOR_RESET}: Run tests in the given order.
EOF
if [ ${#FIO_TEST_MENU_TEXTS[@]} -gt 0 ]; then
    echo -e "                          ${_COLOR_YELLOW}Available Tests (from fio_tests.conf):${_COLOR_RESET}"
    for i in "${!FIO_TEST_MENU_TEXTS[@]}"; do
        printf "                          %s<num> = %d%s: %s\n" "${_COLOR_CYAN}" "$((i+1))" "${_COLOR_RESET}" "${FIO_TEST_MENU_TEXTS[$i]}"
    done
else
    echo -e "                          ${_COLOR_RED}Could not load tests from fio_tests.conf to display here.${_COLOR_RESET}"
fi
cat << EOF
  ${_COLOR_GREEN}-f, --format <0|1>${_COLOR_RESET}
                          Control device formatting (0: No, 1: Yes).
  ${_COLOR_GREEN}-c, --cpu-config-file <path|0>${_COLOR_RESET}
                          Path to a CPU binding config file to override NUMA discovery.
                          ${_COLOR_CYAN}<path> = 0${_COLOR_RESET}: Explicitly disable all CPU binding.
  ${_COLOR_GREEN}-h, --help${_COLOR_RESET}
                          Display this help message and exit.
${_COLOR_YELLOW}Examples:${_COLOR_RESET}
  Display this help message: ./fiotest.sh -h
  Test nvme0n1 without formatting: ./fiotest /dev/nvme0n1 -f 0
  Test all non-system disks with Test Item 7: ./fiotest -d -s 7
  Test multiple disks on items 1, 3, 2 after formatting: ./fiotest /dev/nvme2n1 /dev/nvme3n1 -f 1 -s 1,3,2
  Test nvme1n1 without any CPU core binding: ./fiotest /dev/nvme1n1 -c 0
EOF
}
# --- Configuration Loading ---
create_default_fio_config() {
cat << 'EOF' > "$FIO_TEST_CONFIG_FILE"
# FIO Test Definitions
# Format: "Menu Description"|Log_Filename_Suffix|Base_FIO_Job|rw|bs|iodepth|numjobs
# - 使用竖线 "|" 作为分隔符，以避免描述中的逗号引起解析错误。
"Sequential Preconditioning"|SeqPre|SequentialPrecondition.fio||||
"Sequential Read (1024K, QD32, NJ1)"|SeqRead_1024K_QD32|Common.fio|read|1024K|32|1
"Sequential Write (1024K, QD32, NJ1)"|SeqWrite_1024K_QD32|Common.fio|write|1024K|32|1
"Sequential Write Latency (128K, QD1, NJ1)"|SeqWriteLat_128K_QD1|Common.fio|write|128K|1|1
"Sequential Read Latency (128K, QD1, NJ1)"|SeqReadLat_128K_QD1|Common.fio|read|128K|1|1
"Random Write Preconditioning"|RandPre|RandomPrecondition.fio||||
"Random Read (4K, QD32, NJ8)"|RandRead_4K_QD32|Common.fio|randread|4K|32|8
"Random Write (4K, QD32, NJ8)"|RandWrite_4K_QD32|Common.fio|randwrite|4K|32|8
"Random Write Latency (4K, QD1, NJ1)"|RandWriteLat_4K_QD1|Common.fio|randwrite|4K|1|1
"Random Read Latency (4K, QD1, NJ1)"|RandReadLat_4K_QD1|Common.fio|randread|4K|1|1
EOF
}
load_fio_tests_from_config() {
    if [ ! -f "$FIO_TEST_CONFIG_FILE" ]; then
        log "Configuration file '$FIO_TEST_CONFIG_FILE' not found. Creating a default one." "$_COLOR_YELLOW"
        create_default_fio_config
    fi
    if [ ! -r "$FIO_TEST_CONFIG_FILE" ]; then
        log "Error: Cannot read '$FIO_TEST_CONFIG_FILE'." "$_COLOR_RED"
        return 1
    fi
    local line_num=0
    while IFS= read -r line; do
        line_num=$((line_num + 1))
        # Skip empty lines and comments
        [[ -z "$line" || "$line" =~ ^\s*# ]] && continue
        IFS='|' read -r menu_text log_suffix base_job rw_opt bs_opt iodepth_opt numjobs_opt <<< "$line"
        # Trim whitespace and remove quotes from menu_text
        menu_text="${menu_text//\"/}"
        menu_text="$(echo -e "${menu_text}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')"
        if [ -z "$menu_text" ] || [ -z "$log_suffix" ] || [ -z "$base_job" ]; then
            log "Warning: Malformed line $line_num in '$FIO_TEST_CONFIG_FILE'. Skipping." "$_COLOR_YELLOW"
            continue
        fi
        FIO_TEST_MENU_TEXTS+=("$menu_text")
        FIO_TEST_LOG_SUFFIXES+=("$log_suffix")
        FIO_TEST_BASE_JOBS+=("$base_job")
        FIO_TEST_RW_OPTS+=("$rw_opt")
        FIO_TEST_BS_OPTS+=("$bs_opt")
        FIO_TEST_IODEPTH_OPTS+=("$iodepth_opt")
        FIO_TEST_NUMJOBS_OPTS+=("$numjobs_opt")
    done < "$FIO_TEST_CONFIG_FILE"
    if [ ${#FIO_TEST_MENU_TEXTS[@]} -eq 0 ]; then
        log "Error: No valid test definitions found in '$FIO_TEST_CONFIG_FILE'." "$_COLOR_RED"
        return 1
    fi
    return 0
}
# --- CPU & NUMA Binding Functions ---
DEFAULT_CPU_BINDING_OVERRIDE_FILE="cpu_binding_override.conf"
declare -A device_to_assigned_cores
get_numa_cpus() {
    local node=$1
    if [ -f "/sys/devices/system/node/node${node}/cpulist" ]; then
        cat "/sys/devices/system/node/node${node}/cpulist"
    else
        echo "unknown"
    fi
}
export -f get_numa_cpus
find_pci_device_dir_for_nvme_ctrl() {
    local ctrl_name=$1
    local ctrl_path="/sys/class/nvme/$ctrl_name/device"
    if [ -L "$ctrl_path" ]; then
        readlink -f "$ctrl_path"
    else
        echo ""
    fi
}
export -f find_pci_device_dir_for_nvme_ctrl
generate_nvme_numa_binding_log() {
    local NVME_NUMA_LOGFILE="nvme_numa_binding.log"
    > "$NVME_NUMA_LOGFILE"
    if ! command -v lsblk &> /dev/null; then
        log "Error: lsblk command not found. Cannot determine NUMA bindings." "$_COLOR_RED"
        return 2
    fi
    if ! lsblk -no NAME | grep -q '^nvme'; then
        echo "No NVMe devices found." | tee -a "$NVME_NUMA_LOGFILE"
        return 1
    fi
    local overall_found=false
    # --- Method 1: Check via /sys/class/nvme ---
    detect_via_nvme_class() {
        local method_found=false
        for p in /sys/class/nvme/nvme*; do
            local ctrl_path=$(readlink -f "$p")
            local numa_node_file="$ctrl_path/numa_node"
            if [ -f "$numa_node_file" ]; then
                local node=$(cat "$numa_node_file")
                if [[ "$node" =~ ^[0-9]+$ ]]; then
                    local cpulist=$(get_numa_cpus "$node")
                    echo "$(basename "$p") -> NUMA Node: $node, CPUs: $cpulist"
                    method_found=true
                fi
            fi
        done
        $method_found && return 0 || return 1
    }
    # --- Method 2: Check via /sys/block for namespaces ---
    detect_via_block_namespace() {
        local method_found=false
        for p in /sys/block/nvme*n*; do
            [ ! -d "$p" ] && continue
            local ns_name=$(basename "$p")
            local device_link="${p}/device"
            if [ -L "$device_link" ]; then
                local ctrl_name=$(basename "$(readlink -f "$device_link")")
                local pci_dir=$(find_pci_device_dir_for_nvme_ctrl "$ctrl_name")
                if [ -n "$pci_dir" ]; then
                    local numa_node_file="$pci_dir/numa_node"
                    if [ -f "$numa_node_file" ] && [ -r "$numa_node_file" ]; then
                        local node=$(cat "$numa_node_file")
                        if [ -n "$node" ] && [ "$node" != "-1" ]; then
                            local cpulist=$(get_numa_cpus "$node")
                            echo "$ns_name -> NUMA Node: $node, CPUs: $cpulist"
                            method_found=true
                        fi
                    fi
                fi
            fi
        done
        $method_found && return 0 || return 1
    }
    # --- Method 3: Check via PCI bus uevent ---
    detect_via_pci() {
        local method_found=false
        for n in /sys/class/nvme/nvme*; do
            local ctrl_name=$(basename "$n")
            local ctrl_path=$(readlink -f "$n")
            local uevent_file="$ctrl_path/uevent"
            if [ -f "$uevent_file" ]; then
                local pci_slot=$(grep PCI_SLOT_NAME "$uevent_file" | cut -d= -f2)
                local pci_bus_path="/sys/bus/pci/devices/$pci_slot"
                if [ -d "$pci_bus_path" ]; then
                    local numa_node_file="$pci_bus_path/numa_node"
                    if [ -f "$numa_node_file" ]; then
                        local node=$(cat "$numa_node_file")
                        if [[ "$node" =~ ^[0-9]+$ ]]; then
                            local cpulist=$(get_numa_cpus "$node")
                            echo "$ctrl_name -> NUMA Node: $node, CPUs: $cpulist"
                            method_found=true
                        fi
                    fi
                fi
            fi
        done
        $method_found && return 0 || return 1
    }
    # Run detection methods
    if detect_via_nvme_class >> "$NVME_NUMA_LOGFILE"; then overall_found=true; fi
    if detect_via_block_namespace >> "$NVME_NUMA_LOGFILE"; then overall_found=true; fi
    if detect_via_pci >> "$NVME_NUMA_LOGFILE"; then overall_found=true; fi
    if ! $overall_found; then
        echo "⚠️  Could not determine NUMA node affinity for any NVMe device." | tee "$NVME_NUMA_LOGFILE"
        return 1
    fi
    return 0
}
expand_core_list_to_array() {
    local core_string="$1"
    local expanded_cores=()
    IFS=',' read -ra ranges <<< "$core_string"
    for item in "${ranges[@]}"; do
        if [[ "$item" =~ ^[0-9]+-[0-9]+$ ]]; then
            local start=${item%-*}
            local end=${item#*-}
            if [[ "$start" =~ ^[0-9]+$ && "$end" =~ ^[0-9]+$ && $start -le $end ]]; then
                for (( i=start; i<=end; i++ )); do
                    expanded_cores+=($i)
                done
            else
                log "Warning (expand_core_list_to_array): Invalid range '$item'" >&2
            fi
        elif [[ "$item" =~ ^[0-9]+$ ]]; then
            expanded_cores+=($item)
        elif [ -n "$item" ]; then
            log "Warning (expand_core_list_to_array): Invalid entry '$item'" >&2
        fi
    done
    echo "${expanded_cores[@]}"
}
export -f expand_core_list_to_array
compact_core_array_to_spec() {
    local -a core_int_array=("$@")
    mapfile -t sorted_unique_cores < <(printf "%s\n" "${core_int_array[@]}" | grep -E '^[0-9]+$' | sort -un)
    if [ ${#sorted_unique_cores[@]} -eq 0 ]; then echo ""; return; fi
    local result=""
    local range_start=-1
    local prev_core=-2
    for core_val in "${sorted_unique_cores[@]}"; do
        if (( core_val != prev_core + 1 )); then
            if (( range_start != -1 )); then
                if (( prev_core > range_start )); then
                    result+="${range_start}-${prev_core},"
                else
                    result+="${range_start},"
                fi
            fi
            range_start=$core_val
        fi
        prev_core=$core_val
    done
    if (( range_start != -1 )); then
        if (( prev_core > range_start )); then
            result+="${range_start}-${prev_core},"
        else
            result+="${range_start},"
        fi
    fi
    echo "${result%,}" # Remove trailing comma
}
export -f compact_core_array_to_spec
precalculate_cpu_bindings() {
    log "Starting CPU binding pre-calculation..." "$_COLOR_BLUE"
    device_to_assigned_cores=()
    if [ ! -f "nvme_numa_binding.log" ]; then
        log "Warning: nvme_numa_binding.log not found. Cannot pre-calculate bindings." "$_COLOR_YELLOW"
        return
    fi
    declare -A device_numa_info
    declare -A numa_to_all_cpus
    declare -A numa_to_devices
    for dev_path in "${DEVICES[@]}"; do
        local ctrl_name=$(basename "$dev_path" | sed -E 's/^nvme([0-9]+)n[0-9]+$/nvme\1/')
        local ns_name=$(basename "$dev_path")
        local numa_line=""
        # Prefer namespace-specific line, fallback to controller line
        numa_line=$(grep -m1 -E -i "^${ns_name}([[:space:]]*->[[:space:]]*|[[:space:]]*:)" nvme_numa_binding.log)
        if [ -z "$numa_line" ]; then
            numa_line=$(grep -m1 -E -i "^${ctrl_name}([[:space:]]*->[[:space:]]*|[[:space:]]*:)" nvme_numa_binding.log)
        fi
        if [ -z "$numa_line" ]; then
            log "Warning (precalculate_cpu_bindings): NUMA info for $dev_path not found." >&2
            continue
        fi
        local numa_id=$(echo "$numa_line" | sed -n 's/.*NUMA [Nn]ode:[[:space:]]*\([0-9]\+\).*/\1/p')
        local cpu_spec=$(echo "$numa_line" | sed -n 's/.*CPUs:[[:space:]]*\([0-9,-]\+\)/\1/p' | sed 's/[^0-9,-]//g')
        if [ -z "$numa_id" ] || [ -z "$cpu_spec" ]; then
            log "Warning (precalculate_cpu_bindings): Could not parse NUMA/CPU info for $dev_path" >&2
            continue
        fi
        device_numa_info["$dev_path"]="$numa_id"
        if [ -n "${numa_to_all_cpus["$numa_id"]+_}" ] && [ "${numa_to_all_cpus["$numa_id"]}" != "$cpu_spec" ]; then
            log "Warning (precalculate_cpu_bindings): Inconsistent CPU list for NUMA node $numa_id" >&2
        fi
        numa_to_all_cpus["$numa_id"]="$cpu_spec"
        numa_to_devices["$numa_id"]="${numa_to_devices["$numa_id"]} $dev_path"
    done
    for numa_id in "${!numa_to_all_cpus[@]}"; do
        local cpu_spec_for_node=${numa_to_all_cpus["$numa_id"]}
        read -ra devices_on_node <<< "${numa_to_devices["$numa_id"]}"
        local num_devices_on_node=${#devices_on_node[@]}
        if [ "$num_devices_on_node" -eq 0 ]; then continue; fi
        log "NUMA ${numa_id} (Cores: ${cpu_spec_for_node}) has ${num_devices_on_node} device(s): ${devices_on_node[*]}" "$_COLOR_BLUE"
        read -ra cores_on_node <<< "$(expand_core_list_to_array "$cpu_spec_for_node")"
        local total_cores_on_node=${#cores_on_node[@]}
        if [ "$total_cores_on_node" -eq 0 ]; then
            log "Warning (precalculate_cpu_bindings): NUMA node $numa_id has no usable cores listed." >&2
            for dev in "${devices_on_node[@]}"; do
                device_to_assigned_cores["$dev"]=""
            done
            continue
        fi
        # Distribute cores among devices
        local cores_per_device=$((total_cores_on_node / num_devices_on_node))
        local remainder_cores=$((total_cores_on_node % num_devices_on_node))
        local current_core_index=0
        for i in "${!devices_on_node[@]}"; do
            local current_device=${devices_on_node[$i]}
            local num_cores_to_assign=$cores_per_device
            if [ "$remainder_cores" -gt 0 ]; then
                num_cores_to_assign=$((cores_per_device + 1))
                remainder_cores=$((remainder_cores - 1))
            fi
            local end_index=$((current_core_index + num_cores_to_assign - 1))
            if [ "$end_index" -ge "$total_cores_on_node" ]; then
                end_index=$((total_cores_on_node - 1))
            fi
            local assigned_cores_for_device=()
            if (( current_core_index <= end_index && num_cores_to_assign > 0 )); then
                assigned_cores_for_device=("${cores_on_node[@]:current_core_index:num_cores_to_assign}")
            fi
            local assigned_cpu_spec
            assigned_cpu_spec=$(compact_core_array_to_spec "${assigned_cores_for_device[@]}")
            device_to_assigned_cores["$current_device"]="$assigned_cpu_spec"
            log "  (precalc) Assigned ${current_device}: CPU ${assigned_cpu_spec} (${#assigned_cores_for_device[@]} cores)" "$_COLOR_BLUE"
            current_core_index=$((end_index + 1))
        done
    done
    log "CPU binding pre-calculation finished." "$_COLOR_GREEN"
}
get_binding_cores() {
    local device_path="$1"
    local core_spec=""
    if [ "$DISABLE_CPU_BINDING" == "true" ]; then
        log "[$device_path] (get_binding_cores) CPU binding is globally disabled via -c 0." "$_COLOR_YELLOW" >&2
        echo ""
        return
    fi
    # 1. Highest Priority: User-specified config file via -c
    if [ -n "$USER_SPECIFIED_CPU_CONFIG_FILE" ]; then
        if [ ! -f "$USER_SPECIFIED_CPU_CONFIG_FILE" ]; then
            log "[$device_path]${_COLOR_RED} ERROR: CPU config file '$USER_SPECIFIED_CPU_CONFIG_FILE' not found.${_COLOR_RESET}" >&2
            echo ""
            return
        fi
        local entry
        entry=$(grep -E "^\s*${device_path}\s+" "$USER_SPECIFIED_CPU_CONFIG_FILE" | head -n 1)
        if [ -n "$entry" ]; then
            core_spec=$(echo "$entry" | awk -v dev="$device_path" '$1==dev{$1="";sub(/^[ \t]+/,"");print}')
            if [ -n "$core_spec" ]; then
                log "[$device_path] (get_binding_cores) Using binding from user-specified file: '$core_spec'" "$_COLOR_CYAN" >&2
                echo "$core_spec"
                return
            fi
        fi
        log "[$device_path] (get_binding_cores) No entry found in user-specified file. No binding will be used." "$_COLOR_YELLOW" >&2
        echo ""
        return
    fi
    # 2. Medium Priority: Default override file in script directory
    if [ -f "$DEFAULT_CPU_BINDING_OVERRIDE_FILE" ]; then
        local entry
        entry=$(grep -E "^\s*${device_path}\s+" "$DEFAULT_CPU_BINDING_OVERRIDE_FILE" | head -n 1)
        if [ -n "$entry" ]; then
            core_spec=$(echo "$entry" | awk -v dev="$device_path" '$1==dev{$1="";sub(/^[ \t]+/,"");print}')
            if [ -n "$core_spec" ]; then
                log "[$device_path] (get_binding_cores) Using binding from default override file: '$core_spec'" "$_COLOR_CYAN" >&2
                echo "$core_spec"
                return
            fi
        fi
    fi
    # 3. Low Priority: Dynamic NUMA-based pre-calculation
    if [ -n "${device_to_assigned_cores["$device_path"]+_}" ]; then
        core_spec=${device_to_assigned_cores["$device_path"]}
        if [ -n "$core_spec" ]; then
            sleep 1 # Stagger FIO starts slightly
            log "[$device_path] (get_binding_cores) Using dynamically calculated binding: '$core_spec'" "$_COLOR_BLUE" >&2
            echo "$core_spec"
            return
        else
            log "[$device_path] (get_binding_cores) Dynamic calculation result was empty. No binding." "$_COLOR_YELLOW" >&2
            echo ""
            return
        fi
    fi
    # 4. Fallback: Parse log file again (if dynamic calculation failed or was skipped)
    log "[$device_path] (get_binding_cores) No specific binding found. Falling back to full NUMA node CPUs." "$_COLOR_YELLOW" >&2
    if [ ! -f "nvme_numa_binding.log" ]; then
        log "Warning (get_binding_cores): nvme_numa_binding.log not found for fallback." >&2
        echo ""
        return
    fi
    local ctrl_name=$(basename "$device_path" | sed -E 's/^nvme([0-9]+)n[0-9]+$/nvme\1/')
    local ns_name=$(basename "$device_path")
    local numa_line=""
    numa_line=$(grep -m1 -E -i "^${ns_name}([[:space:]]*->[[:space:]]*|[[:space:]]*:)" nvme_numa_binding.log)
    if [ -z "$numa_line" ]; then
        numa_line=$(grep -m1 -E -i "^${ctrl_name}([[:space:]]*->[[:space:]]*|[[:space:]]*:)" nvme_numa_binding.log)
    fi
    if [ -n "$numa_line" ]; then
        core_spec=$(echo "$numa_line" | sed -n 's/.*CPUs:[[:space:]]*\([0-9,-]\+\)/\1/p' | sed 's/[^0-9,-]//g')
        # Simple cleanup
        core_spec=$(echo "$core_spec" | sed -e 's/--/-/g' -e 's/,,*/,/g' -e 's/^,//' -e 's/,$//')
        echo "$core_spec"
    else
        log "Warning (get_binding_cores): Fallback failed. No info for $ns_name or $ctrl_name." >&2
        echo ""
    fi
}
export -f get_binding_cores
# --- Command Execution ---
execute_cmd() {
    local cmd="$1"
    shift
    local args=("$@")
    if command -v nvme &>/dev/null; then
        nvme "$cmd" "${args[@]}"
    elif [ -x "./nvme" ]; then
        ./nvme "$cmd" "${args[@]}"
    else
        log "Error (execute_cmd): nvme command-line tool not found." "$_COLOR_RED" >&2
        return 1
    fi
}
export -f execute_cmd
format_device() {
    local device_to_format="$1"
    log "Formatting $device_to_format..." "$_COLOR_BLUE"
    execute_cmd format "$device_to_format" -l 0
    if [ $? -ne 0 ]; then
        log "Failed to format $device_to_format. Exiting." "$_COLOR_RED"
        exit 1
    fi
    log "$device_to_format formatted successfully." "$_COLOR_GREEN"
}
export -f format_device
execute_with_taskset() {
    local device="$1"
    shift
    local core_binding_vector
    core_binding_vector=$(get_binding_cores "$device")
    if [ -z "$core_binding_vector" ]; then
        log "[$device] Running FIO without CPU binding." "$_COLOR_YELLOW" >&2
        "$@"
    else
        sleep 2 # Stagger FIO starts slightly for clarity in logs
        log "[$device] Running FIO with CPU binding: $core_binding_vector" "$_COLOR_BLUE" >&2
        taskset -c "$core_binding_vector" "$@"
        local exit_code=$?
        if [ $exit_code -ne 0 ]; then
            log "[$device] (taskset) ${_COLOR_RED}taskset failed with exit code $exit_code. Retrying without binding...${_COLOR_RESET}" >&2
            "$@"
            return $?
        fi
        return $exit_code
    fi
    return $?
}
export -f execute_with_taskset
run_single_fio_test() {
    local device=$1
    local device_dir=$2
    local test_index=$3
    # Note: FIO test definition arrays are now passed via arguments to the parent function
    local base_job="${FIO_TEST_BASE_JOBS[$test_index]}"
    local log_suffix="${FIO_TEST_LOG_SUFFIXES[$test_index]}"
    local rw_opt="${FIO_TEST_RW_OPTS[$test_index]}"
    local bs_opt="${FIO_TEST_BS_OPTS[$test_index]}"
    local iodepth_opt="${FIO_TEST_IODEPTH_OPTS[$test_index]}"
    local numjobs_opt="${FIO_TEST_NUMJOBS_OPTS[$test_index]}"
    local menu_text="${FIO_TEST_MENU_TEXTS[$test_index]}"
    if [ ! -f "$base_job" ]; then
        log "[$device]${_COLOR_RED} Error: FIO job file '$base_job' not found. Skipping test.${_COLOR_RESET}"
        return 1
    fi
    # --- Logic to create categorized subdirectories for detailed logs ---
    local detailed_log_dir="$device_dir" # Default to the base device directory
    case "$menu_text" in
        *[Ss]equential*[Pp]reconditioning*)
            detailed_log_dir="$device_dir/sequential_preconditioning"
            ;;
        *[Rr]andom*[Pp]reconditioning*)
            detailed_log_dir="$device_dir/random_preconditioning"
            ;;
        *[Ss]equential*)
            detailed_log_dir="$device_dir/sequential"
            ;;
        *[Rr]andom*)
            detailed_log_dir="$device_dir/random"
            ;;
    esac
    # Ensure the target directory for detailed logs exists
    mkdir -p "$detailed_log_dir"
    local -a fio_command_args
    fio_command_args+=("$base_job")
    fio_command_args+=(--filename="$device")
    fio_command_args+=(--output="$device_dir/${log_suffix}.out")
    fio_command_args+=(--write_bw_log="$detailed_log_dir/${log_suffix}_bw")
    fio_command_args+=(--write_iops_log="$detailed_log_dir/${log_suffix}_iops")
    fio_command_args+=(--write_lat_log="$detailed_log_dir/${log_suffix}_lat") # This creates _clat,_ slat, _lat
    if [ -n "$rw_opt" ]; then fio_command_args+=(--rw="$rw_opt"); fi
    if [ -n "$bs_opt" ]; then fio_command_args+=(--bs="$bs_opt"); fi
    if [ -n "$iodepth_opt" ]; then fio_command_args+=(--iodepth="$iodepth_opt"); fi
    if [ -n "$numjobs_opt" ]; then fio_command_args+=(--numjobs="$numjobs_opt"); fi
    execute_with_taskset "$device" fio "${fio_command_args[@]}"
}
export -f run_single_fio_test
# --- Function to aggregate FIO logs from multiple jobs ---
aggregate_fio_logs() {
    local log_dir="$1"
    local log_basename="$2" # e.g., "RandRead_4K_QD32_bw"
    local -a files_to_process
    # Use find to locate the numbered log files.
    mapfile -t files_to_process < <(find "$log_dir" -name "${log_basename}.*.log" 2>/dev/null | sort -V)
    # Only perform aggregation if more than one log file part is found.
    if [ "${#files_to_process[@]}" -le 1 ]; then
        return
    fi
    log "[log aggregator] Found ${#files_to_process[@]} log parts for ${log_basename}. Aggregating..." "$_COLOR_BLUE"
    # Use awk to sum the second column (value) grouped by the first column (timestamp).
    # Then pipe to sort to ensure the output file is correctly ordered by timestamp.
    # The output format "timestamp, sum, 0, 0" matches FIO's log format.
    awk -F', ' '{ sums[$1] += $2 } END { for (t in sums) printf "%d, %d, 0, 0\n", t, sums[t] }' "${files_to_process[@]}" | sort -n -t',' -k1 > "${log_dir}/${log_basename}.log"
    if [ $? -eq 0 ]; then
        # If aggregation was successful, remove the original numbered files.
        rm "${files_to_process[@]}"
        log "[log aggregator] Successfully created ${log_basename}.log and cleaned up parts." "$_COLOR_GREEN"
    else
        log "[log aggregator] FAILED to aggregate logs for ${log_basename}." "$_COLOR_RED"
    fi
}
export -f aggregate_fio_logs
# --- Refactored function to run FIO tests in the background ---
run_fio_on_device_background() {
    local device="$1"
    shift
    local -a indices_to_run=("$@")
    # Re-export necessary functions for the subshell environment
    export -f log run_single_fio_test execute_with_taskset get_binding_cores \
           find_pci_device_dir_for_nvme_ctrl get_numa_cpus execute_cmd aggregate_fio_logs
    ( # Start of the subshell
        # Load test definitions. This is crucial for the subshell.
        eval "$FIO_TEST_DEFINITIONS"
        local device_dir="$LOGDIR/$(basename "$device")"
        mkdir -p "$device_dir"
        log_subshell_fail() {
            log "[${device}] FIO Test '$1' FAILED." "$_COLOR_RED"
        }
        for test_idx in "${indices_to_run[@]}"; do
            local log_desc="${FIO_TEST_MENU_TEXTS[$test_idx]}"
            local log_suffix="${FIO_TEST_LOG_SUFFIXES[$test_idx]}"
            local menu_text="${FIO_TEST_MENU_TEXTS[$test_idx]}"
            local numjobs_opt="${FIO_TEST_NUMJOBS_OPTS[$test_idx]}"
            # Determine the subdirectory where logs will be created
            local detailed_log_dir="$device_dir"
            case "$menu_text" in
                *[Ss]equential*[Pp]reconditioning*) detailed_log_dir="$device_dir/sequential_preconditioning" ;;
                *[Rr]andom*[Pp]reconditioning*) detailed_log_dir="$device_dir/random_preconditioning" ;;
                *[Ss]equential*) detailed_log_dir="$device_dir/sequential" ;;
                *[Rr]andom*) detailed_log_dir="$device_dir/random" ;;
            esac
            log "[${device}] Starting test item $((test_idx + 1)): $log_desc..." "$_COLOR_BLUE"
            if run_single_fio_test "$device" "$device_dir" "$test_idx"; then
                # After a successful FIO run, check if we need to aggregate logs.
                # This is only necessary for tests with multiple jobs.
                if [[ -n "$numjobs_opt" && "$numjobs_opt" -gt 1 ]]; then
                    log "[${device}] FIO test completed. Aggregating logs for multi-job test..." "$_COLOR_BLUE"
                    aggregate_fio_logs "$detailed_log_dir" "${log_suffix}_bw"
                    aggregate_fio_logs "$detailed_log_dir" "${log_suffix}_iops"
                    # The --write_lat_log option creates three files with prefixes: _lat,_ slat, _clat
                    aggregate_fio_logs "$detailed_log_dir" "${log_suffix}_lat_lat"
                    aggregate_fio_logs "$detailed_log_dir" "${log_suffix}_lat_slat"
                    aggregate_fio_logs "$detailed_log_dir" "${log_suffix}_lat_clat"
                fi
            else
                log_subshell_fail "$log_desc"
            fi
        done
        if [ ${#indices_to_run[@]} -lt ${#FIO_TEST_MENU_TEXTS[@]} ]; then
            log "[${device}] Selected FIO tests completed." "$_COLOR_GREEN"
        else
            log "[${device}] All FIO tests completed." "$_COLOR_GREEN"
        fi
    ) & FIO_PIDS["$device"]=$!
}
export -f run_fio_on_device_background
handle_interrupt() {
    log "Interrupted. Terminating all background FIO tests..." "$_COLOR_RED"
    if [ ${#FIO_PIDS[@]} -gt 0 ]; then
        log "Sending TERM signal to process groups..." "$_COLOR_YELLOW"
        for key in "${!FIO_PIDS[@]}"; do
            local pid=${FIO_PIDS["$key"]}
            if ps -p "$pid" > /dev/null; then
                # Kill the process group to terminate fio and any children
                kill -- -"$pid" 2>/dev/null || kill "$pid" 2>/dev/null
            fi
        done
        sleep 5
        log "Sending KILL signal to any remaining processes..." "$_COLOR_YELLOW"
        for key in "${!FIO_PIDS[@]}"; do
            local pid=${FIO_PIDS["$key"]}
            if ps -p "$pid" > /dev/null; then
                kill -9 -- -"$pid" 2>/dev/null || kill -9 "$pid" 2>/dev/null
            fi
        done
    else
        log "No background FIO PIDs to kill." "$_COLOR_YELLOW"
    fi
    log "Also killing any remaining fio processes by name..." "$_COLOR_YELLOW"
    pkill -f "fio.*--filename=/dev/nvme"
    log "Interrupt handling complete. Exiting." "$_COLOR_RED"
    sleep 1
    exit 130
}
trap handle_interrupt INT TERM

parse_and_display_summary() {
    # --- Local Helper Function for Bandwidth Extraction/Conversion ---
    _get_mbs() {
        local summary_line="$1"
        # --- Method 1: Directly extract the (MB/s) value from modern FIO output.
        local direct_mbs; direct_mbs=$(echo "$summary_line" | grep -o -E '\([0-9.]+MB/s\)' | grep -o -E '[0-9.]+')
        if [[ -n "$direct_mbs" ]]; then
            printf "%.2f\n" "$direct_mbs"
            return
        fi
        # --- Method 2: Fallback for older FIO versions or unusual output.
        local raw_bw; raw_bw=$(echo "$summary_line" | sed -n 's/.*BW=\([^ ]\+\).*/\1/p')
        if [[ -z "$raw_bw" || "$raw_bw" == "N/A" ]]; then echo "N/A"; return; fi
        local val; val=$(echo "$raw_bw" | grep -o -E '[0-9.]+')
        local unit_prefix; unit_prefix=$(echo "$raw_bw" | tr '[:lower:]' '[:upper:]' | grep -o -E 'GIB|MIB|KIB|G|M|K|B')
        if [[ -z "$val" ]]; then echo "N/A"; return; fi
        local mbs="0.00"
        # 1 MiB = 1.048576 MB
        case "$unit_prefix" in
            GIB|G) mbs=$(bc -l <<< "$val * 1024 * 1.048576") ;;
            MIB|M) mbs=$(bc -l <<< "$val * 1.048576") ;;
            KIB|K) mbs=$(bc -l <<< "$val * 1024 / 1000000") ;;
            B)     mbs=$(bc -l <<< "$val / 1000000") ;;
            *)     echo "$raw_bw"; return ;;
        esac
        printf "%.2f\n" "$mbs"
    }
    echo -e "\n${_COLOR_CYAN}======================================================================${_COLOR_RESET}"
    echo -e "${_COLOR_CYAN}                      FIO Performance Test Result                   ${_COLOR_RESET}"
    echo -e "${_COLOR_CYAN}======================================================================${_COLOR_RESET}"
    local indices_to_parse=()
    if [ ${#SELECTED_FIO_STEP_INDICES[@]} -gt 0 ]; then
        indices_to_parse=("${SELECTED_FIO_STEP_INDICES[@]}")
    else
        indices_to_parse=($(seq 0 $(( ${#FIO_TEST_MENU_TEXTS[@]} - 1 )) ))
    fi
    # --- Pass 1: Collect all data ---
    declare -A results
    log "Parsing results for summary table..."
    for device in "${DEVICES[@]}"; do
        for test_idx in "${indices_to_parse[@]}"; do
            local device_name; device_name=$(basename "$device")
            local log_suffix="${FIO_TEST_LOG_SUFFIXES[$test_idx]}"
            local output_file="$LOGDIR/$device_name/${log_suffix}.out"
            local iops="N/A" bw_mbs="N/A" lat_str="N/A" 
            if [ -f "$output_file" ]; then
                local summary_line; summary_line=$(grep -m1 -E ' read: IOPS| write: IOPS' "$output_file")
                iops=$(echo "$summary_line" | sed -n 's/.*IOPS=\([0-9kK.]\+\).*/\1/p')

                
                local lat_line; lat_line=$(grep -m1 -E "^\s+lat \(" "$output_file")
                local lat_avg; lat_avg=$(echo "$lat_line" | sed -n 's/.*avg=\([0-9.]\+\).*/\1/p')

                bw_mbs=$(_get_mbs "$summary_line")
                lat_str="${lat_avg:-N/A}"
                if [[ "$lat_str" != "N/A" ]]; then
                    # Check the unit from the same line we got the value from
                    if echo "$lat_line" | grep -q "(usec)"; then lat_str+=" us";
                    elif echo "$lat_line" | grep -q "(msec)"; then lat_str+=" ms";
                    elif echo "$lat_line" | grep -q "(nsec)"; then lat_str+=" ns"; fi
                fi
            fi
            results["$device,$test_idx,iops"]="${iops:-N/A}"
            results["$device,$test_idx,bw"]="${bw_mbs}"
            results["$device,$test_idx,lat"]="${lat_str}" 
        done
    done
    # --- Pass 2: Build and print the transposed table  ---
    local header_printf_format="%-30s"
    local data_printf_format="%-30s"
    local header_values=("Test Metric")
    local separator_line
    separator_line=$(printf "%-30s" "" | tr ' ' '-')
    for device in "${DEVICES[@]}"; do
        header_values+=("$(basename "$device")")
        header_printf_format+=" | %-12s"
        data_printf_format+=" | %-12s"
        separator_line+=$(printf "+%-12s" "" | tr ' ' '-')
    done
    # Print Header
    printf "${_COLOR_YELLOW}$header_printf_format\n${_COLOR_RESET}" "${header_values[@]}"
    echo "$separator_line"
    # Print Data Rows, grouped by test
    for test_idx in "${indices_to_parse[@]}"; do
        local suffix="${FIO_TEST_LOG_SUFFIXES[$test_idx]}"
        local iops_row_data=("${suffix}_IOPS")
        for device in "${DEVICES[@]}"; do iops_row_data+=("${results["$device,$test_idx,iops"]}"); done
        printf "$data_printf_format\n" "${iops_row_data[@]}"
        local bw_row_data=("${suffix}_BW(MB/s)")
        for device in "${DEVICES[@]}"; do bw_row_data+=("${results["$device,$test_idx,bw"]}"); done
        printf "$data_printf_format\n" "${bw_row_data[@]}"
        local clat_row_data=("${suffix}_LAT") # Variable name here is just for display, doesn't matter
        for device in "${DEVICES[@]}"; do clat_row_data+=("${results["$device,$test_idx,lat"]}"); done
        printf "$data_printf_format\n" "${clat_row_data[@]}"
        if [[ "$test_idx" != "${indices_to_parse[-1]}" ]]; then
            echo "$separator_line" | tr '-' ' ' | sed 's/|/+/g'
        fi
    done
    echo "$separator_line"
}
export -f parse_and_display_summary
# --- #################### ---
# --- MAIN SCRIPT LOGIC    ---
# --- #################### ---
# --- Pre-flight Checks and Initialization ---
load_fio_tests_from_config || exit 1
FIO_TEST_DEFINITIONS=$(declare -p FIO_TEST_MENU_TEXTS FIO_TEST_LOG_SUFFIXES FIO_TEST_BASE_JOBS FIO_TEST_RW_OPTS FIO_TEST_BS_OPTS FIO_TEST_IODEPTH_OPTS FIO_TEST_NUMJOBS_OPTS)
export FIO_TEST_DEFINITIONS
# Early exit for help request
for arg in "$@"; do
    if [[ "$arg" == "-h" || "$arg" == "--help" ]]; then
        display_usage
        exit 0
    fi
done
# --- Argument Parsing ---
temp_args_for_getopt=()
original_args=("$@")
arg_idx=0
while [ $arg_idx -lt ${#original_args[@]} ]; do
    arg="${original_args[$arg_idx]}"
    case "$arg" in
        -s|--select-fio-step|-f|--format|-c|--cpu-config-file)
            temp_args_for_getopt+=("$arg")
            arg_idx=$((arg_idx+1))
            if [ $arg_idx -lt ${#original_args[@]} ]; then
                temp_args_for_getopt+=("${original_args[$arg_idx]}")
            else
                log "Error: Option $arg requires an argument." "$_COLOR_RED"
                exit 1
            fi
            ;;
        -d|--all-non-system-nvme)
            temp_args_for_getopt+=("$arg")
            ;;
        /dev/nvme*n[0-9]*)
            if [ -b "$arg" ]; then
                CLI_DEVICE_ARGS+=("$arg")
            else
                log "Warning: Device '$arg' (from command line) is not a valid block device." "$_COLOR_YELLOW"
            fi
            ;;
        *)
            temp_args_for_getopt+=("$arg")
            ;;
    esac
    arg_idx=$((arg_idx+1))
done
TEMP_ARGS=$(getopt -o df:s:c:h --long all-non-system-nvme,format:,select-fio-step:,cpu-config-file:,help -n "$0" -- "${temp_args_for_getopt[@]}")
if [ $? -ne 0 ]; then
    echo -e "${_COLOR_RED}Parameter parsing error. Use -h or --help for usage.${_COLOR_RESET}" >&2
    exit 1
fi
eval set -- "$TEMP_ARGS"
unset TEMP_ARGS
while true; do
    case "$1" in
        -s|--select-fio-step)
            step_choice_str="$2"
            max_step_num=${#FIO_TEST_MENU_TEXTS[@]}
            if [[ "$step_choice_str" == "0" ]]; then
                RUN_FULL_FIO_VIA_S0=true
            else
                IFS=',' read -ra step_nums_arr <<< "$step_choice_str"
                for step_num in "${step_nums_arr[@]}"; do
                    [[ -z "$step_num" ]] && continue
                    if ! [[ "$step_num" =~ ^[0-9]+$ ]]; then
                        log "Error: Invalid non-numeric value '$step_num' provided to -s option." "$_COLOR_RED"
                        exit 1
                    fi
                    if [ "$step_num" -ge 1 ] && [ "$step_num" -le "$max_step_num" ]; then
                        SELECTED_FIO_STEP_INDICES+=($(( step_num - 1 )))
                    else
                        log "Error: Invalid FIO step number '$step_num'. Must be between 1 and $max_step_num, or 0 for all." "$_COLOR_RED"
                        exit 1
                    fi
                done
                if [ ${#SELECTED_FIO_STEP_INDICES[@]} -eq 0 ]; then
                    log "Error: No valid FIO steps were selected with the -s option." "$_COLOR_RED"
                    exit 1
                fi
            fi
            shift 2
            ;;
        -d|--all-non-system-nvme)
            AUTO_SELECT_ALL_NON_SYSTEM_NVME=true
            shift
            ;;
        -f|--format)
            PERFORM_FORMATTING="$2"
            if [[ "$PERFORM_FORMATTING" != "0" && "$PERFORM_FORMATTING" != "1" ]]; then
                log "Error: Invalid argument for -f/--format. Must be 0 or 1." "$_COLOR_RED"
                exit 1
            fi
            shift 2
            ;;
        -c|--cpu-config-file)
            if [[ "$2" == "0" ]]; then
                DISABLE_CPU_BINDING="true"
                log "CPU binding explicitly disabled via '-c 0' option." "$_COLOR_CYAN"
            else
                USER_SPECIFIED_CPU_CONFIG_FILE="$2"
                if [ ! -f "$USER_SPECIFIED_CPU_CONFIG_FILE" ]; then
                    log "Error: CPU config file '$USER_SPECIFIED_CPU_CONFIG_FILE' not found." "$_COLOR_RED"
                    exit 1
                fi
            fi
            shift 2
            ;;
        -h|--help)
            display_usage
            exit 0
            ;;
        --)
            shift
            break
            ;;
        *)
            log "Internal getopt error: '$1'." "$_COLOR_RED"
            exit 1
            ;;
    esac
done
if [ "$#" -gt 0 ]; then
    log "Error: Unrecognized arguments after options: '$@'." "$_COLOR_RED"
    exit 1
fi
# Check if user specified tests in a non-sequential order and ask for confirmation
if [ ${#SELECTED_FIO_STEP_INDICES[@]} -gt 1 ]; then
    is_out_of_order=false
    last_index=-1
    for current_index in "${SELECTED_FIO_STEP_INDICES[@]}"; do
        if [ "$last_index" -ne -1 ] && [ "$current_index" -lt "$last_index" ]; then
            is_out_of_order=true
            break
        fi
        last_index=$current_index
    done
    if "$is_out_of_order"; then
        sequence_str=""
        user_numbers_arr=()
        for index in "${SELECTED_FIO_STEP_INDICES[@]}"; do
            sequence_str+="'${FIO_TEST_MENU_TEXTS[$index]}' -> "
            user_numbers_arr+=($((index + 1)))
        done
        sequence_str="${sequence_str% -> }"
        user_numbers_str=$(IFS=,; echo "${user_numbers_arr[*]}")
        log "Non-sequential test order detected." "$_COLOR_YELLOW"
        log "Your specified test order is ${user_numbers_str}, which corresponds to:" "$_COLOR_CYAN"
        log "${sequence_str}" "$_COLOR_CYAN"
        read -p "Are you sure you want to run tests in this order? (y=continue, n=exit): " -n 1 -r REPLY
        echo
        if [[ ! "$REPLY" =~ ^[Yy]$ ]]; then
            log "User canceled the test. Exiting." "$_COLOR_RED"
            exit 0
        fi
        log "User confirmed non-sequential order. Proceeding." "$_COLOR_GREEN"
    fi
fi
log "FIO Test Script Initializing..." "$_COLOR_CYAN"
log "Log directory: $LOGDIR" "$_COLOR_BLUE"
if [ "$EUID" -ne 0 ]; then
    log "Error: This script must be run as root." "$_COLOR_RED"
    exit 1
fi
# --- Device Selection Logic ---
populate_and_display_selectable_nvme_devices() {
    AVAILABLE_DEVICE_PATHS_FOR_SELECTION=()
    AVAILABLE_DEVICES_DISPLAY_INFO=()
    # Find the root partition's parent disk to exclude it
    local root_dev_info
    root_dev_info=$(findmnt -n -o SOURCE / | grep -E '^/dev/nvme[0-9]+n[0-9]+(p[0-9]+)?$' || true)
    local system_nvme_parent_disk=""
    if [[ "$root_dev_info" =~ ^/dev/nvme([0-9]+n[0-9]+)(p[0-9]+)?$ ]]; then
        system_nvme_parent_disk="/dev/nvme${BASH_REMATCH[1]}"
    fi
    # Use lsblk with parseable output format
    local all_nvme_info
    all_nvme_info=$(lsblk -P -o NAME,TYPE,SIZE,FSTYPE,MOUNTPOINT,PKNAME --paths -e 7,11)
    if [ -z "$all_nvme_info" ]; then
        echo "No NVMe devices found."
        return
    fi
    declare -A all_disk_paths disk_info_map disk_to_partition_details
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        declare -A devinfo
        eval "$(echo "$line" | awk '{ n = split($0, p, /[[:space:]]+/); for (i=1; i<=n; i++) { if (match(p[i], /^([^=]+)="(.*)"/, m)) { gsub(/^"|"$/, "", m[2]); gsub(/\\/, "\\\\", m[2]); print "devinfo[\"" m[1] "\"]=\"" m[2] "\";" } } }')"
        local dev_name="${devinfo[NAME]}"
        local dev_type="${devinfo[TYPE]}"
        local dev_size="${devinfo[SIZE]}"
        local dev_fstype="${devinfo[FSTYPE]}"
        local dev_mountpoint="${devinfo[MOUNTPOINT]}"
        local pkname="${devinfo[PKNAME]}"
        if [[ -z "$dev_name" || ! "$dev_name" =~ ^/dev/nvme[0-9]+n[0-9]+(p[0-9]+)?$ ]]; then
            continue
        fi
        if [[ "$dev_type" == "disk" && "$dev_name" =~ ^/dev/nvme[0-9]+n[0-9]+$ ]]; then
            all_disk_paths["$dev_name"]=true
            disk_info_map["$dev_name"]="${dev_type}|${dev_size}|${dev_fstype}|${dev_mountpoint}"
        elif [[ "$dev_type" == "part" && "$pkname" =~ ^/dev/nvme[0-9]+n[0-9]+$ ]]; then
            DEVICE_HAS_PARTITIONS["$pkname"]=true
            local part_display_info="${dev_name} (${dev_size}"
            [[ -n "$dev_type" ]] && part_display_info+=" $dev_type"
            [[ -n "$dev_fstype" ]] && part_display_info+=" $dev_fstype"
            [[ -n "$dev_mountpoint" ]] && part_display_info+=" mounted at $dev_mountpoint"
            part_display_info+=")"
            if [ -z "${disk_to_partition_details["$pkname"]}" ]; then
                disk_to_partition_details["$pkname"]="$part_display_info"
            else
                disk_to_partition_details["$pkname"]="${disk_to_partition_details["$pkname"]};${part_display_info}"
            fi
        fi
    done <<< "$all_nvme_info"
    echo -e "${_COLOR_CYAN}--- Available Non-System NVMe Devices ---${_COLOR_RESET}"
    local counter=1
    local sorted_disks
    sorted_disks=($(for k in "${!all_disk_paths[@]}"; do echo "$k"; done | sort -V))
    for disk_path in "${sorted_disks[@]}"; do
        if [[ -n "$system_nvme_parent_disk" && "$disk_path" == "$system_nvme_parent_disk" ]]; then
            continue
        fi
        local disk_info_str="${disk_info_map["$disk_path"]}"
        IFS='|' read -r dev_type dev_size dev_fstype dev_mountpoint <<< "$disk_info_str"
        local info_display="(${dev_size}"
        [[ -n "$dev_type" ]] && info_display+=" $dev_type"
        [[ -n "$dev_fstype" ]] && info_display+=" $dev_fstype"
        [[ -n "$dev_mountpoint" ]] && info_display+=" mounted at $dev_mountpoint"
        info_display+=")"
        printf " %2d) %s %s\n" "$counter" "$disk_path" "$info_display"
        if [ "${DEVICE_HAS_PARTITIONS["$disk_path"]}" == "true" ]; then
            IFS=';' read -ra pa <<< "${disk_to_partition_details["$disk_path"]}"
            for pe in "${pa[@]}"; do
                echo "    -> $pe"
            done
        fi
        AVAILABLE_DEVICE_PATHS_FOR_SELECTION+=("$disk_path")
        AVAILABLE_DEVICES_DISPLAY_INFO+=("$disk_path $info_display")
        counter=$((counter + 1))
    done
    echo -e "${_COLOR_CYAN}---------------------------------${_COLOR_RESET}"
    if [ "$counter" -eq 1 ]; then
        echo "No non-system NVMe disks found."
    fi
    # Temporarily export partition details for the warning message later
    declare -p disk_to_partition_details > "$LOGDIR/disk_to_partition_details_export.tmp"
}
export -f populate_and_display_selectable_nvme_devices
if "$AUTO_SELECT_ALL_NON_SYSTEM_NVME"; then
    TEST_TYPE="2" # Multi-disk mode
    if [ ${#CLI_DEVICE_ARGS[@]} -gt 0 ]; then
        log "Warning: Positional device arguments are ignored when -d is used." "$_COLOR_YELLOW"
    fi
    log "Mode: Auto-selecting all non-system NVMe disks." "$_COLOR_GREEN"
    populate_and_display_selectable_nvme_devices
    DEVICES=("${AVAILABLE_DEVICE_PATHS_FOR_SELECTION[@]}")
    if [ ${#DEVICES[@]} -eq 0 ]; then
        log "Error: No devices found for auto-selection." "$_COLOR_RED"
        exit 1
    fi
    log "Auto-selected devices: ${DEVICES[*]}" "$_COLOR_GREEN"
elif [ ${#CLI_DEVICE_ARGS[@]} -gt 0 ]; then
    DEVICES=("${CLI_DEVICE_ARGS[@]}")
    if [ "${#DEVICES[@]}" -eq 1 ]; then
        TEST_TYPE="1" # Single-disk mode
        log "Single disk test selected via command line: ${DEVICES[*]}" "$_COLOR_GREEN"
    else
        TEST_TYPE="2" # Multi-disk mode
        log "Multi-disk test selected via command line: ${DEVICES[*]}" "$_COLOR_GREEN"
    fi
    populate_and_display_selectable_nvme_devices # Display for context
else
    log "Starting interactive device selection..." "$_COLOR_CYAN"
    populate_and_display_selectable_nvme_devices
    if [ ${#AVAILABLE_DEVICE_PATHS_FOR_SELECTION[@]} -eq 0 ]; then
        log "Error: No devices available for selection." "$_COLOR_RED"
        exit 1
    fi
    while true; do
        read -p "Select test type (1: Single Disk, 2: Multi Disk): " TEST_TYPE
        if [[ "$TEST_TYPE" == "1" || "$TEST_TYPE" == "2" ]]; then
            break
        else
            echo -e "${_COLOR_RED}Invalid input. Please enter 1 or 2.${_COLOR_RESET}"
        fi
    done
    if [ "$TEST_TYPE" == "1" ]; then
        while true; do
            read -p "Enter device number, full path, or 'q' to quit: " sel
            if [[ "$sel" =~ ^[Qq]$ ]]; then exit 0; fi
            if [[ "$sel" =~ ^[0-9]+$ ]]; then
                num=$((sel - 1))
                if [ "$num" -ge 0 ] && [ "$num" -lt ${#AVAILABLE_DEVICE_PATHS_FOR_SELECTION[@]} ]; then
                    DEVICES=("${AVAILABLE_DEVICE_PATHS_FOR_SELECTION[$num]}")
                    break
                else
                    echo -e "${_COLOR_RED}Invalid number.${_COLOR_RESET}"
                fi
            elif [ -b "$sel" ]; then
                if [[ "$sel" =~ p[0-9]+$ ]]; then
                    echo -e "${_COLOR_RED}Cannot select a partition directly. Please select the parent disk.${_COLOR_RESET}"
                    continue
                fi
                DEVICES=("$sel")
                break
            else
                echo -e "${_COLOR_RED}Invalid input. Please enter a number from the list, a device path, or 'q'.${_COLOR_RESET}"
            fi
        done
        log "Selected device: ${DEVICES[*]}" "$_COLOR_GREEN"
    elif [ "$TEST_TYPE" == "2" ]; then
        while true; do
            read -p "Enter device numbers (space-separated, e.g., '1 3') or 'q' to quit: " sel
            if [[ "$sel" =~ ^[Qq]$ ]]; then exit 0; fi
            if [ -n "$sel" ]; then
                read -ra nums <<< "$sel"
                paths=()
                invalid=false
                for n_str in "${nums[@]}"; do
                    if [[ "$n_str" =~ ^[0-9]+$ ]]; then
                        num=$((n_str - 1))
                        if [ "$num" -ge 0 ] && [ "$num" -lt ${#AVAILABLE_DEVICE_PATHS_FOR_SELECTION[@]} ]; then
                            paths+=("${AVAILABLE_DEVICE_PATHS_FOR_SELECTION[$num]}")
                        else
                            echo -e "${_COLOR_RED}Invalid device number: '$n_str'.${_COLOR_RESET}"; invalid=true; break
                        fi
                    else
                        echo -e "${_COLOR_RED}Invalid input: '$n_str' is not a number.${_COLOR_RESET}"; invalid=true; break
                    fi
                done
                if ! "$invalid" && [ ${#paths[@]} -gt 0 ]; then
                    mapfile -t DEVICES < <(printf "%s\n" "${paths[@]}" | sort -u)
                    break
                else
                    echo -e "${_COLOR_RED}No valid devices were selected. Please try again.${_COLOR_RESET}"
                fi
            else
                echo -e "${_COLOR_RED}No input provided.${_COLOR_RESET}"
            fi
        done
        log "Selected devices: ${DEVICES[*]}" "$_COLOR_GREEN"
    fi
fi
if [ ${#DEVICES[@]} -eq 0 ]; then
    log "Error: No devices were selected for testing. Exiting." "$_COLOR_RED"
    exit 1
fi

if "$AUTO_SELECT_ALL_NON_SYSTEM_NVME" && ! ${RUN_FULL_FIO_VIA_S0} && [ ${#SELECTED_FIO_STEP_INDICES[@]} -eq 0 ]; then
    read -p "Auto-test all selected disks with the full FIO suite? (y/n) " -n 1 -r REPLY
    echo
    if [[ "$REPLY" =~ ^[Yy]$ ]]; then
        log "User confirmed running the full FIO suite." "$_COLOR_GREEN"
        RUN_FULL_FIO_VIA_S0=true # User agreed to run all tests
    else
        log "Proceeding to manual test selection..." "$_COLOR_CYAN"
    fi
fi

# --- Pre-Test Warnings and Confirmations ---
declare -A disk_to_partition_details
if [ -f "$LOGDIR/disk_to_partition_details_export.tmp" ]; then
    source "$LOGDIR/disk_to_partition_details_export.tmp"
    rm "$LOGDIR/disk_to_partition_details_export.tmp"
fi
warning_message_parts=()
for dev_path in "${DEVICES[@]}"; do
    if [ "${DEVICE_HAS_PARTITIONS["$dev_path"]}" == "true" ]; then
        partitions_str=""
        if [ -n "${disk_to_partition_details["$dev_path"]}" ]; then
            partitions_str=" ($(IFS=,; echo "${disk_to_partition_details["$dev_path"]//;/}"))"
        fi
        warning_message_parts+=("${dev_path}${partitions_str}")
    fi
done
if [ ${#warning_message_parts[@]} -gt 0 ]; then
    log "\n${_COLOR_RED}WARNING: The following selected disks have existing partitions:\n  $(IFS=$'\n  '; echo "${warning_message_parts[*]}")\nTesting or formatting these disks may result in DATA LOSS.${_COLOR_RESET}\n"
    read -p "Do you want to continue with these partitioned disks? (y/n) " -n 1 -r REPLY
    echo
    if [[ ! "$REPLY" =~ ^[Yy]$ ]]; then
        log "User cancelled due to partition warning. Exiting." "$_COLOR_RED"
        exit 0
    fi
fi
if [ ${#SELECTED_FIO_STEP_INDICES[@]} -eq 0 ] && ! "$RUN_FULL_FIO_VIA_S0"; then
    max_step_num=${#FIO_TEST_MENU_TEXTS[@]}
    while true; do
        log "FIO Test Steps:" "$_COLOR_BLUE"
        log "  0) Run ALL"
        for i in "${!FIO_TEST_MENU_TEXTS[@]}"; do
            log "  $((i+1))) ${FIO_TEST_MENU_TEXTS[$i]}"
        done
        read -p "Enter step number(s) to run (e.g., '1,3,2' or '0' for all): " sel
        if [[ "$sel" == "0" ]]; then
            RUN_FULL_FIO_VIA_S0=true
            break
        fi
        temp_indices=()
        invalid_found=false
        IFS=',' read -ra step_nums_arr <<< "$sel"
        for step_num in "${step_nums_arr[@]}"; do
            [[ -z "$step_num" ]] && continue
            if ! [[ "$step_num" =~ ^[0-9]+$ ]]; then
                echo -e "${_COLOR_RED}Invalid input: '$step_num' is not a number.${_COLOR_RESET}"; invalid_found=true; break
            fi
            if [ "$step_num" -ge 1 ] && [ "$step_num" -le "$max_step_num" ]; then
                temp_indices+=($(( step_num - 1 )))
            else
                echo -e "${_COLOR_RED}Invalid step number: '$step_num'. Must be between 1 and $max_step_num.${_COLOR_RESET}"; invalid_found=true; break
            fi
        done
        if ! "$invalid_found" && [ ${#temp_indices[@]} -gt 0 ]; then
            SELECTED_FIO_STEP_INDICES=("${temp_indices[@]}")
            log "Selected steps: ${sel}" "$_COLOR_BLUE"
            break
        fi
    done
fi
if [ -z "$PERFORM_FORMATTING" ]; then
    read -p "Format selected devices (${DEVICES[*]}) before testing? THIS IS DESTRUCTIVE. (y/n): " -n 1 -r REPLY
    echo
    if [[ "$REPLY" =~ ^[Yy]$ ]]; then
        PERFORM_FORMATTING="1"
    else
        PERFORM_FORMATTING="0"
    fi
fi
# --- Test Execution Phase ---
log "Script Setup & Prerequisite Checks..." "$_COLOR_BLUE"
if [ -f "./nvme" ]; then chmod 777 ./nvme 2>/dev/null; fi
find . -maxdepth 1 -type f \( -name "*.sh" -o -name "*.fio" \) -exec chmod 777 {} + 2>/dev/null
if [ "$DISABLE_CPU_BINDING" != "true" ]; then
    log "Generating NVMe NUMA binding info..." "$_COLOR_BLUE"
    generate_nvme_numa_binding_log
    if [ $? -ne 0 ]; then
        log "NUMA binding info generation may have failed. CPU affinity might not be optimal." "$_COLOR_RED"
    else
        log "NUMA binding info generated successfully." "$_COLOR_GREEN"
    fi
else
    log "Skipping NUMA binding info generation as CPU binding is disabled." "$_COLOR_YELLOW"
fi
log "Retrieving NVMe device list for information log..."
tf=$(mktemp)
if execute_cmd list > "$tf" 2>&1; then
    cat "$tf" > "$LOGDIR/information.log"
else
    cat "$tf" > "$LOGDIR/information.log"
    log "Failed to get full device list from nvme-cli." "$_COLOR_RED"
fi
rm -f "$tf"
if [ "$DISABLE_CPU_BINDING" != "true" ] && [ -z "$USER_SPECIFIED_CPU_CONFIG_FILE" ]; then
    precalculate_cpu_bindings
elif [ -n "$USER_SPECIFIED_CPU_CONFIG_FILE" ]; then
    log "Using user-specified CPU config file; skipping dynamic CPU allocation." "$_COLOR_CYAN"
else # This case handles when DISABLE_CPU_BINDING is true
    log "Skipping dynamic CPU allocation as binding is disabled." "$_COLOR_CYAN"
fi
if [ "$PERFORM_FORMATTING" == "1" ]; then
    log "Formatting devices as requested..." "$_COLOR_BLUE"
    for device_to_format in "${DEVICES[@]}"; do
        format_device "$device_to_format"
    done
    log "Formatting complete. Waiting 30 seconds for devices to settle..." "$_COLOR_GREEN"
    sleep 30
else
    log "Skipping device formatting." "$_COLOR_YELLOW"
fi
# Determine which test indices to run
indices_to_run=()
if [ ${#SELECTED_FIO_STEP_INDICES[@]} -gt 0 ]; then
    indices_to_run=("${SELECTED_FIO_STEP_INDICES[@]}")
else
    indices_to_run=($(seq 0 $(( ${#FIO_TEST_MENU_TEXTS[@]} - 1 )) ))
fi
log "Starting FIO testing process in parallel..." "$_COLOR_CYAN"
for device in "${DEVICES[@]}"; do
    # Call the refactored background function with explicit parameters
    run_fio_on_device_background "$device" "${indices_to_run[@]}"
done
# --- Wait for all background jobs and report results ---
ALL_COMPLETED_SUCCESSFULLY=true
log "Waiting for all background FIO tests to complete..." "$_COLOR_BLUE"
for dev_pid_key in "${!FIO_PIDS[@]}"; do
    PID=${FIO_PIDS["$dev_pid_key"]}
    if [ -z "$PID" ]; then
        log "Warning: PID for device '$dev_pid_key' was not found." "$_COLOR_YELLOW"
        continue
    fi
    log "Waiting for FIO on '$dev_pid_key' (PID $PID)..."
    wait "$PID"
    EXIT_STATUS=$?
    if [ $EXIT_STATUS -ne 0 ]; then
        log "FIO on '$dev_pid_key' FAILED with exit code: $EXIT_STATUS." "$_COLOR_RED"
        ALL_COMPLETED_SUCCESSFULLY=false
    else
        log "FIO on '$dev_pid_key' finished successfully." "$_COLOR_GREEN"
    fi
done
if "$ALL_COMPLETED_SUCCESSFULLY"; then
    log "All FIO tests completed successfully. Results are in: $LOGDIR." "$_COLOR_GREEN"
else
    log "One or more FIO tests FAILED. Please check the logs in '$LOGDIR'." "$_COLOR_YELLOW"
fi
log "Script finished." "$_COLOR_CYAN"
parse_and_display_summary